{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ca3bd17d-aec3-4848-ac82-def6e2d6fa18",
   "metadata": {},
   "source": [
    "# Examples of Structured Data Extraction in LlamaIndex\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/structured_outputs/structured_outputs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "If you haven't yet read our [structured data extraction tutorial](../../understanding/extraction/index.md), we recommend starting there. This notebook demonstrates some of the techniques introduced in the tutorial.\n",
    "\n",
    "We start with the simple syntax around LLMs, then move on to how to use it with higher-level modules like a query engine and agent.\n",
    "\n",
    "A lot of the underlying behavior around structured outputs is powered by our Pydantic Program modules. Check out our [in-depth structured outputs guide](https://docs.llamaindex.ai/en/stable/module_guides/querying/structured_outputs/) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6625f456-bf9c-4c89-948c-488909591855",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57ced509-b1fd-4413-9ac2-662fcfab1074",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4o\")\n",
    "embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72a15e8d-45dc-4aef-93a6-8d2e97604d1e",
   "metadata": {},
   "source": [
    "## 1. Simple Structured Extraction\n",
    "\n",
    "You can convert any LLM to a \"structured LLM\" by attaching an output class to it through `as_structured_llm`.\n",
    "\n",
    "Here we pass a simple `Album` class which contains a list of songs. We can then use the normal LLM endpoints like chat/complete.\n",
    "\n",
    "**NOTE**: async is supported but streaming is coming soon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63045118-74fb-4ec6-bb99-b66184eba017",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class Song(BaseModel):\n",
    "    \"\"\"Data model for a song.\"\"\"\n",
    "\n",
    "    title: str\n",
    "    length_seconds: int\n",
    "\n",
    "\n",
    "class Album(BaseModel):\n",
    "    \"\"\"Data model for an album.\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "    songs: List[Song]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5929511-67c3-411a-925c-5155286d10cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "sllm = llm.as_structured_llm(output_cls=Album)\n",
    "input_msg = ChatMessage.from_str(\"Generate an example album from The Shining\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a8fd46-dbc9-4973-bf70-aa0f28a7fa2d",
   "metadata": {},
   "source": [
    "#### Sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4671ba4b-d21f-4a4d-8aa8-555a679248c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = sllm.chat([input_msg])\n",
    "# get actual object\n",
    "output_obj = output.raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c58e29dc-b6e5-4fdb-83c4-ccb011dde2fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 720}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 540}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 300}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 480}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 540}, {\"title\": \"Home\", \"length_seconds\": 180}, {\"title\": \"Midnight, the Stars and You\", \"length_seconds\": 180}, {\"title\": \"It's All Forgotten Now\", \"length_seconds\": 150}, {\"title\": \"Masquerade\", \"length_seconds\": 180}]}\n",
      "name='The Shining: Original Soundtrack' artist='Various Artists' songs=[Song(title='Main Title', length_seconds=180), Song(title='Rocky Mountains', length_seconds=210), Song(title='Lontano', length_seconds=720), Song(title='Music for Strings, Percussion and Celesta', length_seconds=540), Song(title='Utrenja (Excerpt)', length_seconds=300), Song(title='The Awakening of Jacob', length_seconds=480), Song(title='De Natura Sonoris No. 2', length_seconds=540), Song(title='Home', length_seconds=180), Song(title='Midnight, the Stars and You', length_seconds=180), Song(title=\"It's All Forgotten Now\", length_seconds=150), Song(title='Masquerade', length_seconds=180)]\n"
     ]
    }
   ],
   "source": [
    "print(str(output))\n",
    "print(output_obj)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a93195eb-b62b-43eb-acad-0652f745bc64",
   "metadata": {},
   "source": [
    "#### Async"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c4e582-5aa8-42e8-8d27-d9b18b8452d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title (The Shining)\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 240}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 300}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 180}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 150}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 270}, {\"title\": \"Home\", \"length_seconds\": 200}, {\"title\": \"Heartbeats and Worry\", \"length_seconds\": 160}, {\"title\": \"The Overlook\", \"length_seconds\": 220}]}\n"
     ]
    }
   ],
   "source": [
    "output = await sllm.achat([input_msg])\n",
    "# get actual object\n",
    "output_obj = output.raw\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6608650-c3a9-4c09-92af-9c61376d291c",
   "metadata": {},
   "source": [
    "#### Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "403ab86e-b1ca-4f4c-bcb3-4c5173ed17f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'artist': 'Various Artists',\n",
      " 'name': 'The Shining: Original Soundtrack',\n",
      " 'songs': [{'length_seconds': 180, 'title': 'Main Title'},\n",
      "           {'length_seconds': 210, 'title': 'Rocky Mountains'},\n",
      "           {'length_seconds': 240, 'title': 'Lontano'},\n",
      "           {'length_seconds': 540,\n",
      "            'title': 'Music for Strings, Percussion and Celesta'},\n",
      "           {'length_seconds': 300, 'title': 'Utrenja (Excerpt)'},\n",
      "           {'length_seconds': 360, 'title': 'The Awakening of Jacob'},\n",
      "           {'length_seconds': 420, 'title': 'De Natura Sonoris No. 2'},\n",
      "           {'length_seconds': 180, 'title': 'Home'},\n",
      "           {'length_seconds': 180, 'title': 'Midnight, the Stars and You'},\n",
      "           {'length_seconds': 150, 'title': \"It's All Forgotten Now\"},\n",
      "           {'length_seconds': 120, 'title': 'Masquerade'}]}\n",
      "assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title (The Shining)\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 240}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 300}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 180}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 150}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 270}, {\"title\": \"Home\", \"length_seconds\": 200}, {\"title\": \"Heartbeats and Worry\", \"length_seconds\": 160}, {\"title\": \"The Overlook\", \"length_seconds\": 220}]}\n"
     ]
    }
   ],
   "source": [
    "from IPython.display import clear_output\n",
    "from pprint import pprint\n",
    "\n",
    "stream_output = sllm.stream_chat([input_msg])\n",
    "for partial_output in stream_output:\n",
    "    clear_output(wait=True)\n",
    "    pprint(partial_output.raw.dict())\n",
    "\n",
    "output_obj = partial_output.raw\n",
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "900037f6-3fa9-4f55-bcfc-985131517f36",
   "metadata": {},
   "source": [
    "#### Async Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c0025b-b657-42c8-a873-f51929431b8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'artist': 'Various Artists',\n",
      " 'name': 'The Shining: Original Soundtrack',\n",
      " 'songs': [{'length_seconds': 180, 'title': 'Main Title'},\n",
      "           {'length_seconds': 210, 'title': 'Rocky Mountains'},\n",
      "           {'length_seconds': 720, 'title': 'Lontano'},\n",
      "           {'length_seconds': 540,\n",
      "            'title': 'Music for Strings, Percussion and Celesta'},\n",
      "           {'length_seconds': 300, 'title': 'Utrenja (Excerpt)'},\n",
      "           {'length_seconds': 480, 'title': 'The Awakening of Jacob'},\n",
      "           {'length_seconds': 540, 'title': 'De Natura Sonoris No. 2'},\n",
      "           {'length_seconds': 180, 'title': 'Home'},\n",
      "           {'length_seconds': 180, 'title': 'Midnight, the Stars and You'},\n",
      "           {'length_seconds': 180, 'title': \"It's All Forgotten Now\"},\n",
      "           {'length_seconds': 180, 'title': 'Masquerade'}]}\n"
     ]
    }
   ],
   "source": [
    "from IPython.display import clear_output\n",
    "from pprint import pprint\n",
    "\n",
    "stream_output = await sllm.astream_chat([input_msg])\n",
    "async for partial_output in stream_output:\n",
    "    clear_output(wait=True)\n",
    "    pprint(partial_output.raw.dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d1b4bc9-56ee-4bf7-a63c-50ce0cadeeae",
   "metadata": {},
   "source": [
    "### 1.b Use the `structured_predict` Function\n",
    "\n",
    "Instead of explicitly doing `llm.as_structured_llm(...)`, every LLM class has a `structured_predict` function which allows you to more easily call the LLM with a prompt template + template variables to return a strutured output in one line of code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2488b69e-c04f-48fe-afea-7224dc79efa9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Album(name='Songs of Middle-earth', artist='Various Artists', songs=[Song(title='The Shire', length_seconds=180), Song(title='The Fellowship', length_seconds=240), Song(title=\"Gollum's Theme\", length_seconds=200), Song(title=\"Rohan's Call\", length_seconds=220), Song(title=\"The Battle of Helm's Deep\", length_seconds=300), Song(title='Lothlórien', length_seconds=210), Song(title='The Return of the King', length_seconds=250), Song(title='Into the West', length_seconds=260)])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# use query pipelines\n",
    "from llama_index.core.prompts import ChatPromptTemplate\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "chat_prompt_tmpl = ChatPromptTemplate(\n",
    "    message_templates=[\n",
    "        ChatMessage.from_str(\n",
    "            \"Generate an example album from {movie_name}\", role=\"user\"\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4o\")\n",
    "album = llm.structured_predict(\n",
    "    Album, chat_prompt_tmpl, movie_name=\"Lord of the Rings\"\n",
    ")\n",
    "album"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae13bf7e-888a-4305-8ed8-d22f130ab6fe",
   "metadata": {},
   "source": [
    "## 2. Plug into RAG Pipeline\n",
    "\n",
    "You can also plug this into a RAG pipeline. Below we show structured extraction from an Apple 10K report."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c72151ca-b471-41ad-a762-f483ca341d9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "!wget \"https://s2.q4cdn.com/470004039/files/doc_financials/2021/q4/_10-K-2021-(As-Filed).pdf\" -O data/apple_2021_10k.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1f81633-cdcb-4e8d-a43f-43b699a0cf01",
   "metadata": {},
   "source": [
    "#### Option 1: Use LlamaParse\n",
    "\n",
    "You will need an account at https://cloud.llamaindex.ai/ and an API Key to use LlamaParse, our document parser for 10K filings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2745476-02be-4f2a-86ff-5386aae1260d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id cac11eca-7e00-452f-93f6-19c861b4c130\n"
     ]
    }
   ],
   "source": [
    "from llama_parse import LlamaParse\n",
    "\n",
    "# os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"llx-...\"\n",
    "orig_docs = LlamaParse(result_type=\"text\").load_data(\n",
    "    \"./data/apple_2021_10k.pdf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9509b85e-0091-4fc0-8649-5a3d83d6dcf0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                              UNITED STATES\n",
      "                                    SECURITIES AND EXCHANGE COMMISSION\n",
      "                                                          Washington, D.C. 20549\n",
      "\n",
      "                                                               FORM 10-K\n",
      "(Mark One)\n",
      "        ☒ ANNUAL REPORT PURSUANT TO SECTION 13 OR 15(d) OF THE SECURITIES EXCHANGE ACT OF 1934\n",
      "                                                  For the fiscal year ended September 25, 2021\n",
      "                                                                            or\n",
      "     ☐ TRANSITION REPORT PURSUANT TO SECTION 13 OR 15(d) OF THE SECURITIES EXCHANGE ACT OF 1934\n",
      "                                                For the transition period from               to          .\n",
      "                                                       Commission File Number: 001-36743\n",
      "\n",
      "                                                               Apple Inc.\n",
      "                                              (Exact name of Registrant as specified in its charter)\n",
      "\n",
      "                            California                                                                          94-2404110\n",
      "                     (State or other jurisdiction                                                    (I.R.S. Employer Identification No.)\n",
      "                  of incorporation or organization)\n",
      "\n",
      "                     One Apple Park Way\n",
      "                     Cupertino, California                                                                         95014\n",
      "              (Address of principal executive offices)                                                            (Zip Code)\n",
      "                                                                    (408) 996-1010\n",
      "                                                 (Registrant’s telephone number, including area code)\n",
      "\n",
      "                                               Securities registered pursuant to Section 12(b) of the Act:\n",
      "\n",
      "                                                                         Trading\n",
      "                         Title of each class                            symbol(s)               Name of each exchange on which registered\n",
      "      Common Stock, $0.00001 par value per share                         AAPL                       The Nasdaq Stock Market LLC\n",
      "                    1.000% Notes due 2022                                   —                       The Nasdaq Stock Market LLC\n",
      "                    1.375% Notes due 2024                                   —                       The Nasdaq Stock Market LLC\n",
      "                    0.000% Notes due 2025                                   —                       The Nasdaq Stock Market LLC\n",
      "                    0.875% Notes due 2025                                   —                       The Nasdaq Stock Market LLC\n",
      "                    1.625% Notes due 2026                                   —                       The Nasdaq Stock Market LLC\n",
      "                    2.000% Notes due 2027                                   —                       The Nasdaq Stock Market LLC\n",
      "                    1.375% Notes due 2029                                   —                       The Nasdaq Stock Market LLC\n",
      "                    3.050% Notes due 2029                                   —                       The Nasdaq Stock Market LLC\n",
      "                    0.500% Notes due 2031                                   —                       The Nasdaq Stock Market LLC\n",
      "                    3.600% Notes due 2042                                   —                       The Nasdaq Stock Market LLC\n",
      "\n",
      "                                           Securities registered pursuant to Section 12(g) of the Act: None\n",
      "\n",
      "Indicate by check mark if the Registrant is a well-known seasoned issuer, as defined in Rule 405 of the Securities Act.\n",
      "                                                                    Yes ☒      No ☐\n",
      "Indicate by check mark if the Registrant is not required to file reports pursuant to Section 13 or Section 15(d) of the Act.\n",
      "                                                                    Yes ☐      No ☒\n"
     ]
    }
   ],
   "source": [
    "from copy import deepcopy\n",
    "from llama_index.core.schema import TextNode\n",
    "\n",
    "\n",
    "def get_page_nodes(docs, separator=\"\\n---\\n\"):\n",
    "    \"\"\"Split each document into page node, by separator.\"\"\"\n",
    "    nodes = []\n",
    "    for doc in docs:\n",
    "        doc_chunks = doc.text.split(separator)\n",
    "        for doc_chunk in doc_chunks:\n",
    "            node = TextNode(\n",
    "                text=doc_chunk,\n",
    "                metadata=deepcopy(doc.metadata),\n",
    "            )\n",
    "            nodes.append(node)\n",
    "\n",
    "    return nodes\n",
    "\n",
    "\n",
    "docs = get_page_nodes(orig_docs)\n",
    "print(docs[0].get_content())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ae6c400-e661-4a5c-8ed0-1f31bf8325cb",
   "metadata": {},
   "source": [
    "#### Option 2: Use SimpleDirectoryReader\n",
    "\n",
    "You can also choose to use the free PDF parser bundled into our `SimpleDirectoryReader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee757df9-0491-46ba-b369-d8cdf1b7d74f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # OPTION 2: Use SimpleDirectoryReader\n",
    "# from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# reader = SimpleDirectoryReader(input_files=[\"apple_2021_10k.pdf\"])\n",
    "# docs = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f3e9304-1377-46e4-9f23-939907c6198d",
   "metadata": {},
   "source": [
    "#### Build RAG Pipeline, Define Structured Output Schema\n",
    "\n",
    "We build a RAG pipeline with our trusty VectorStoreIndex and reranker module. We then define the output as a Pydantic model. This allows us to create a structured LLM with the output class attached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e318de5-7c7b-49a2-8de6-140f16b07d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "# skip chunking since we're doing page-level chunking\n",
    "index = VectorStoreIndex(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6df8be76-53f6-4e81-8b22-7281e0ac35eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor.flag_embedding_reranker import (\n",
    "    FlagEmbeddingReranker,\n",
    ")\n",
    "\n",
    "reranker = FlagEmbeddingReranker(\n",
    "    top_n=5,\n",
    "    model=\"BAAI/bge-reranker-large\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e52dfc-31a5-41bf-972c-284c8a3263d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Output(BaseModel):\n",
    "    \"\"\"Output containing the response, page numbers, and confidence.\"\"\"\n",
    "\n",
    "    response: str = Field(..., description=\"The answer to the question.\")\n",
    "    page_numbers: List[int] = Field(\n",
    "        ...,\n",
    "        description=\"The page numbers of the sources used to answer this question. Do not include a page number if the context is irrelevant.\",\n",
    "    )\n",
    "    confidence: float = Field(\n",
    "        ...,\n",
    "        description=\"Confidence value between 0-1 of the correctness of the result.\",\n",
    "    )\n",
    "    confidence_explanation: str = Field(\n",
    "        ..., description=\"Explanation for the confidence score\"\n",
    "    )\n",
    "\n",
    "\n",
    "sllm = llm.as_structured_llm(output_cls=Output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f28e413-c9ee-44dd-88c3-0f8c387341e5",
   "metadata": {},
   "source": [
    "#### Run Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c2d822-50cb-4eca-9c14-661c47bb70b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=5,\n",
    "    node_postprocessors=[reranker],\n",
    "    llm=sllm,\n",
    "    response_mode=\"tree_summarize\",  # you can also select other modes like `compact`, `refine`\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc211576-7101-4d1d-a874-eaf1ba5ad7a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"response\": \"In 2021, the net sales for each product category were as follows: iPhone: $191,973 million, Mac: $35,190 million, iPad: $31,862 million, Wearables, Home and Accessories: $38,367 million, and Services: $68,425 million.\", \"page_numbers\": [21], \"confidence\": 1.0, \"confidence_explanation\": \"The figures are directly taken from the provided data, ensuring high accuracy.\"}\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Net sales for each product category in 2021\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "340bc8a8-7460-4683-bf2e-518c35ba68af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'response': 'In 2021, the net sales for each product category were as follows: iPhone: $191,973 million, Mac: $35,190 million, iPad: $31,862 million, Wearables, Home and Accessories: $38,367 million, and Services: $68,425 million.',\n",
       " 'page_numbers': [21],\n",
       " 'confidence': 1.0,\n",
       " 'confidence_explanation': 'The figures are directly taken from the provided data, ensuring high accuracy.'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.response.dict()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v3",
   "language": "python",
   "name": "llama_index_v3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
